import numpy as np
import h5py
import sys
from math import pi, sqrt
def writeinput(fname,k,t,Efield,E,H,rho0):
    fil=h5py.File(fname,'w')
    fil.create_dataset('kpoints',data=k);
    fil.create_dataset('time',data=t);
    fil.create_dataset('Efield',data=Efield);
    fil.create_dataset('EnergyBands',data=E);
    fil.create_dataset('InteractionHamiltonian_real',data=np.real(H));
    fil.create_dataset('InteractionHamiltonian_imag',data=np.imag(H));
    fil.create_dataset('rho0_real',data=np.real(rho0));
    fil.create_dataset('rho0_imag',data=np.imag(rho0));
    fil.close()
    return 0

def input1():
    nt   = 100
    tend = 10

    nkx  = 33
    nky  = 2
    nb   = 2

    freq = 1.0;
    Egap = 0.5;
        
    t  = np.linspace(0,tend,nt);
    kx = np.linspace(0,0.4,nkx);
    ky = np.linspace(0,0.4,nky);
    
    k1,k2 = np.meshgrid(kx,ky);
    
    k  = np.vstack((k1.flatten(),k2.flatten())).T
    nk = k.shape[0];

    dt = t[1]-t[0]


    Efield = 0.2*np.exp(-(t-t[-1]/2)**2/2/16)*np.cos(freq*(t-t[-1]/2));
           
    E      = np.zeros((nk,nb));
    H      = np.zeros((nk,nb,nb));

    E[:,0] =  -0.1*np.sum(k**2,axis=1);
    E[:,1] =   np.sum(k**2,axis=1)+Egap;

    H[:,0,1] =  sqrt(8.0)* pi;
    H[:,1,0] =  np.conj(H[:,0,1])
    
    rho0 = np.zeros((nk,nb,nb),dtype=np.complex);

    rho0[:,1,1]=1.0

    return k,t,Efield,E,H,rho0

def uniformly_discretized_Afield(t,Efield,numlevels=10):
    Efield=Efield.flatten()
    Afield=scipy.integrate.cumtrapz(Efield,x=t)
    Afield=np.hstack((0,Afield))
    Alevels = np.linspace(np.min(Afield),np.max(Afield),numlevels)
    T=[]
    Afieldi=[]
    for e in Alevels:
        A = Afield[0:-1] - e
        B = Afield[1:]   - e
        i = np.argwhere(A*B<0).flatten()
        for j in i:
            ti = t[j]+(e-Afield[j])/(Afield[j+1]-Afield[j])*(t[j+1]-t[j])
            T.append(ti)
            Afieldi.append(e)
            
    T=np.array(T)
    i=np.argsort(T)
    T=T[i]
    Afieldi=np.array(Afieldi)
    Afieldi=Afieldi[i]
    return T,Afieldi,Afield
    
if __name__=='__main__':
    import scipy.integrate
    
    k,t,Efield,E,H,rho0=input1()
    writeinput('input.h5',k,t,Efield,E,H,rho0)
    
    T,Afieldi,Afield=uniformly_discretized_Afield(t,Efield,numlevels=10)
    
    
        
    
    
